package com.feriki.msync.process.sink;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.feriki.msync.constants.PolicyType;
import com.feriki.msync.config.ComplexConnectConfig;
import com.feriki.msync.constants.ConnectType;
import com.feriki.msync.meta.MilvusCollectionMeta;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.tuple.Pair;

import java.math.BigDecimal;
import java.math.BigInteger;
import java.sql.*;
import java.util.*;
import java.util.Date;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**clickhouseSink仅支持追加数据
 * @author luozhida
 * @create 2024-12-05
 */
@Slf4j
public class ClickhouseSinkHandler extends SinkHandler {
    private final String INSERT_SQL_TEMPLATE = "INSERT INTO %s(%s) VALUES(%s)";

    private ComplexConnectConfig connectConfig;
    private int clickhouseNumberTypeIsNullFill;
    private List<String> clickhouseFloat32CompressTo16BitEmbeddingFieldList;
    private Connection connection;
    private AtomicBoolean prepared = new AtomicBoolean(false);
    //确保key有序
    private Map<String, Pair<Integer,String>> columnNameTypeMetaMap = new LinkedHashMap<>();
    private String insertSql;
    private Pattern dateNumberPattern;

    public ClickhouseSinkHandler(String tableName, ComplexConnectConfig connectConfig, int clickhouseNumberTypeIsNullFill,String clickhouseFloat32CompressTo16BitEmbeddingFields) {
        super(tableName, ConnectType.CH, PolicyType.APPEND);
        this.connectConfig = connectConfig;
        this.clickhouseNumberTypeIsNullFill = clickhouseNumberTypeIsNullFill;
        this.connection = createConnection();
        if(clickhouseFloat32CompressTo16BitEmbeddingFields!=null){
            this.clickhouseFloat32CompressTo16BitEmbeddingFieldList = Arrays.asList(clickhouseFloat32CompressTo16BitEmbeddingFields.split(","));
        }
        log.info("落地至clickhouse暂仅支持"+PolicyType.APPEND+"策略!");
    }

    @Override
    public boolean prepare(String source, MilvusCollectionMeta sourceMeta) {
        if (!prepared.get()) {
            String queryColumnSql = "select * from " + getName() + " where 1=2";
            Statement statement = null;
            ResultSet rs = null;
            try {
                statement = connection.createStatement();

                rs = statement.executeQuery(queryColumnSql);
                ResultSetMetaData rsMetaData = rs.getMetaData();
                columnNameTypeMetaMap.clear();
                for (int i = 0, len = rsMetaData.getColumnCount(); i < len; i++) {
                    columnNameTypeMetaMap.put(rsMetaData.getColumnName(i + 1), Pair.of(rsMetaData.getColumnType(i + 1),rsMetaData.getColumnTypeName(i + 1)));
                }
            } catch (SQLException e) {
                if(e.getMessage().contains("doesn't exist")||e.getMessage().contains("UNKNOWN_TABLE")){
                    log.error("clickhouse表{}不存在,请自行创建,且字段命名需要与源数据的字段名一致!",getName());
                }else{
                    log.error("clickhouse表{}元数据信息查询失败",getName(),e);
                }
                return false;
            } finally {
                if (null != rs) {
                    try {
                        rs.close();
                    } catch (SQLException unused) {
                    }
                }

                if (null != statement) {
                    try {
                        statement.close();
                    } catch (SQLException unused) {
                    }
                }
            }
            String columnNames  = StringUtils.join(columnNameTypeMetaMap.keySet(),",");
            String holds = StringUtils.repeat("?,",columnNameTypeMetaMap.keySet().size());
            holds = holds.substring(0,holds.length()-1);
            insertSql = String.format(INSERT_SQL_TEMPLATE,getName(),columnNames,holds);
            log.info("执行的SQL是:"+insertSql);
            prepared.compareAndSet(false, true);
        }
        return true;
    }


    /**
    *  经试验，不支持
    * @create 2024-12-16
    */
    private String getSqlColumnHolds() {
        List<String> sqlColumnHoldList = new ArrayList<String>();
        for(String column:columnNameTypeMetaMap.keySet()){
            if(clickhouseFloat32CompressTo16BitEmbeddingFieldList!=null&&
                    clickhouseFloat32CompressTo16BitEmbeddingFieldList.contains(column)){
                sqlColumnHoldList.add("arrayMap(x -> reinterpretAsFloat32(bitAnd(reinterpretAsUInt32(x), 4294901760)), ?)");
            }else{
                sqlColumnHoldList.add("?");
            }
        }
        return StringUtils.join(sqlColumnHoldList,",");
    }

    @Override
    public long handler(String source, List<JSONObject> sourceJsonObjects) {
        PreparedStatement preparedStatement = null;
        boolean handled = false;
        try {
            connection = validDBConn(connection);

            connection.setAutoCommit(false);
            preparedStatement = connection.prepareStatement(this.insertSql);

         //   preparedStatement.executeQuery("SYSTEM STOP MERGES");
            for(JSONObject sourceJsonObject:sourceJsonObjects){
                fillPreparedStatement(preparedStatement,sourceJsonObject);
                preparedStatement.addBatch();
            }

            preparedStatement.executeBatch();
            connection.commit();
            handled = true;
        } catch (Exception e) {
            if (e instanceof SQLException) {
                try {
                    connection.rollback();//使用同一个connection执行多个sql,如果第一个失败但没有回滚，之后的sql会抛异常：“当前事务被终止，命令被忽略，直到事务块结束”
                } catch (SQLException ex) {
                    log.error("连接回滚失败，将重新申请连接", e);
                    this.destroy();
                    connection = createConnection();
                }
            }
            //打日志后丢弃
            log.error(String.format("当前从源【%s】批量同步到目标clickhouse表【%s】%s条记录失败", source, getName(), sourceJsonObjects.size()),e);
            return 0;
        }finally {
            if(preparedStatement!=null){
                try {
                    /*if(handled){
                        preparedStatement.executeQuery("SYSTEM START MERGES");
                    }*/
                    preparedStatement.close();
                } catch (SQLException e) {
                    //ignore
                }
            }
        }
        log.info(String.format("当前成功从源【%s】批量同步到目标clickhouse表【%s】共%s条记录", source, getName(), sourceJsonObjects.size()));
        return sourceJsonObjects.size();
    }

    @Override
    public void destroy() {
        if (connection != null) {
            try {
                connection.close();
            } catch (SQLException e) {
                //ignore
            }
        }

    }

    private void fillPreparedStatement(PreparedStatement preparedStatement, JSONObject sourceJsonObject)
            throws Exception {
        int i = 0;
        for (Map.Entry<String,Pair<Integer,String>> entry:columnNameTypeMetaMap.entrySet()) {
            fillPreparedStatementColumnType(preparedStatement, i, entry.getKey(),
                    entry.getValue(), sourceJsonObject);
            i++;
        }
    }

    private void fillPreparedStatementColumnType(PreparedStatement preparedStatement,
                                                                int columnIndex,
                                                                String columnName,Pair<Integer,String> columnTypePair,
                                                                JSONObject sourceJsonObject) throws Exception {
        String sourceStrValue = sourceJsonObject.getString(columnName);
        switch (columnTypePair.getLeft()) {
            case Types.CHAR:
            case Types.NCHAR:
            case Types.CLOB:
            case Types.NCLOB:
            case Types.VARCHAR:
            case Types.LONGVARCHAR:
            case Types.NVARCHAR:
            case Types.LONGNVARCHAR:
                preparedStatement.setString(columnIndex + 1,sourceStrValue);
                break;
            case Types.TINYINT:
            case Types.SMALLINT:
            case Types.INTEGER:
            case Types.BOOLEAN:
                BigInteger bigInteger = sourceJsonObject.getBigInteger(columnName);
                int intValue = bigInteger!=null?bigInteger.intValue():clickhouseNumberTypeIsNullFill;
                preparedStatement.setInt(columnIndex + 1, intValue);
                break;
            case Types.BIGINT:
                Long longTemp = sourceJsonObject.getLong(columnName);
                long longValue = longTemp!=null?longTemp:clickhouseNumberTypeIsNullFill;
                preparedStatement.setLong(columnIndex + 1,longValue );
                break;
            case Types.DECIMAL:
                BigDecimal bigDecimalTemp = sourceJsonObject.getBigDecimal(columnName);
                BigDecimal bigDecimalValue = bigDecimalTemp!=null?bigDecimalTemp:new BigDecimal(clickhouseNumberTypeIsNullFill);
                preparedStatement.setBigDecimal(columnIndex + 1, bigDecimalValue);
                break;
            case Types.REAL:
            case Types.FLOAT:
                Float floatTemp = sourceJsonObject.getFloat(columnName);
                float floatValue = floatTemp!=null?floatTemp:clickhouseNumberTypeIsNullFill;
                preparedStatement.setFloat(columnIndex + 1, floatValue);
                break;
            case Types.DOUBLE:
                Double doubleTemp = sourceJsonObject.getDouble(columnName);
                double doubleValue = doubleTemp!=null?doubleTemp:clickhouseNumberTypeIsNullFill;
                preparedStatement.setDouble(columnIndex + 1, doubleValue);
                break;
            case Types.DATE:
                java.sql.Date sqlDate = null;
                if (columnTypePair.getRight()
                        .equalsIgnoreCase("year")) {
                    if (sourceJsonObject.getBigInteger(columnName) == null) {
                        preparedStatement.setString(columnIndex + 1, null);
                    } else {
                        preparedStatement.setInt(columnIndex + 1, sourceJsonObject.getBigInteger(columnName).intValue());
                    }
                }else{
                    try {
                        Date dateValue = sourceJsonObject.getDate(columnName);
                        if (null != dateValue) {
                            sqlDate = new java.sql.Date(dateValue.getTime());
                        }
                    }catch (Exception ex){
                        log.warn("字段【{}】DATE类型转换出错:",columnName);
                    }
                    preparedStatement.setDate(columnIndex + 1, sqlDate);
                }
                break;
            case Types.TIME:
                java.sql.Time sqlTime = null;
                try {
                    Date dateValue = sourceJsonObject.getDate(columnName);
                    if (null != dateValue) {
                        sqlTime = new java.sql.Time(dateValue.getTime());
                    }
                }catch (Exception ex){
                    log.warn("字段【{}】TIME类型转换出错:",columnName);
                }
                preparedStatement.setTime(columnIndex + 1, sqlTime);
                break;
            case Types.TIMESTAMP:
                java.sql.Timestamp sqlTimestamp = null;
                // JAVA TIMESTAMP 类型入参必须是 "2017-07-12 14:39:00.123566" 格式
                if(dateNumberPattern==null){
                    String pattern = "^\\d+-\\d+-\\d+ \\d+:\\d+:\\d+.\\d+";
                    dateNumberPattern = Pattern.compile(pattern);
                }
                Matcher m = dateNumberPattern.matcher(sourceStrValue);
                if (m.matches()) {
                    sqlTimestamp = Timestamp.valueOf(sourceStrValue);
                    preparedStatement.setTimestamp(columnIndex + 1, sqlTimestamp);
                }else{
                    try {
                        Date dateValue = sourceJsonObject.getDate(columnName);
                        if (null != dateValue) {
                            sqlTimestamp = new java.sql.Timestamp(dateValue.getTime());
                        }
                    }catch (Exception ex){
                        log.warn("字段【{}】TIMESTAMP类型转换出错:",columnName,ex);
                    }
                    preparedStatement.setTimestamp(columnIndex + 1, sqlTimestamp);
                }
                break;
            case Types.BINARY:
            case Types.BLOB:
            case Types.VARBINARY:
            case Types.LONGVARBINARY:
                preparedStatement.setBytes(columnIndex + 1, sourceJsonObject.getBytes(columnName));
                break;
            case Types.BIT:
                String bitValue="0";
                if(StringUtils.isNotBlank(sourceStrValue)){
                    bitValue=sourceStrValue;
                }
                preparedStatement.setString(columnIndex + 1, bitValue);
                break;
            case Types.OTHER:
                if(!columnTypePair.getRight().startsWith("Tuple")) {
                    preparedStatement.setString(columnIndex + 1, sourceStrValue);
                }else{
                    preparedStatement.setObject(columnIndex + 1,sourceJsonObject.get(columnName));
                }
                break;
            case Types.ARRAY:
                JSONArray values = sourceJsonObject.getJSONArray(columnName);
                String componentTypeName = values.get(0).getClass().getSimpleName();
                Object[] convertArray = convertNumberArray(values,componentTypeName);
                Array array = connection.createArrayOf(componentTypeName, convertArray);
                preparedStatement.setArray(columnIndex + 1, array);
                break;
            default:
                preparedStatement.setObject(columnIndex + 1,sourceJsonObject.get(columnName));
                break;
        }
    }

    private Connection validDBConn(Connection connection){
        try {
            if(connection==null || connection.isClosed() || !connection.isValid(5)){
                return createConnection();
            }
        } catch (SQLException e) {
            throw new RuntimeException("验证连接有效性出错",e);
        }
        return connection;
    }

    private Connection createConnection(){
        Properties prop = new Properties();
        prop.put("user", this.connectConfig.getUser());
        if(this.connectConfig.getPassword()!=null){
            prop.put("password", this.connectConfig.getPassword());
        }

        try {
            Class.forName("ru.yandex.clickhouse.ClickHouseDriver");
            DriverManager.setLoginTimeout(60);
            return DriverManager.getConnection(this.connectConfig.getJdbcUrl(), prop);
        } catch (ClassNotFoundException | SQLException e) {
            throw new RuntimeException("初始化clickhouse连接失败",e);
        }
    }

    private Object[] convertNumberArray(JSONArray values,String toType){
        //适配一下向量数据(数值型数组)
        List<String> originalValues = values.toJavaList(String.class);
        String lowerCaseType = toType.toLowerCase();
        if(lowerCaseType.contains("float") || lowerCaseType.contains("double")){
            Float[] floatList = new Float[originalValues.size()];
            for (int i = 0; i < originalValues.size(); i++) {
                floatList[i] = Float.valueOf(originalValues.get(i));
            }
            return floatList;
        }else if(lowerCaseType.contains("int")){
            Integer[] intList = new Integer[originalValues.size()];
            for (int i = 0; i < originalValues.size(); i++) {
                intList[i] = Integer.valueOf(originalValues.get(i));
            }
            return intList;
        }else{
            return originalValues.toArray();
        }
    }
}